{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use Amazon Sagemaker Distributed Model Parallel to Launch a BERT Training Job with Model Parallelization\n",
    "\n",
    "Sagemaker distributed model parallel (SMP) is a model parallelism library for training large deep learning models that were previously difficult to train due to GPU memory limitations. SMP automatically and efficiently splits a model across multiple GPUs and instances and coordinates model training, allowing you to increase prediction accuracy by creating larger models with more parameters.\n",
    "\n",
    "Use this notebook to configure SMP to train a model using PyTorch (version 1.6.0) and the [Amazon SageMaker Python SDK](https://sagemaker.readthedocs.io/en/stable/overview.html#train-a-model-with-the-sagemaker-python-sdk).\n",
    "\n",
    "In this notebook, you will use a BERT example training script with SMP.\n",
    "The example script is based on [Nvidia Deep Learning Examples](https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/LanguageModeling/BERT) and requires you to download the datasets and upload them to Amazon Simple Storage Service (Amazon S3) as explained in the instructions below. This is a large dataset, and so depending on your connection speed, this process can take hours to complete. \n",
    "\n",
    "This notebook depends on the following files. You can find all files in the [bert directory](https://github.com/aws/amazon-sagemaker-examples/tree/master/training/distributed_training/pytorch/model_parallel/bert) in the model parllel section of the Amazon SageMaker Examples notebooks repo.\n",
    "\n",
    "* `bert_example/sagemaker_smp_pretrain.py`: This is an entrypoint script that is passed to the Pytorch estimator in the notebook instructions. This script is responsible for end to end training of the BERT model with SMP. The script has additional comments at places where the SMP API is used.\n",
    "\n",
    "* `bert_example/modeling.py`: This contains the model definition for the BERT model.\n",
    "\n",
    "* `bert_example/bert_config.json`: This allows for additional configuration of the model and is used by `modeling.py`. Additional configuration includes dropout probabilities, pooler and encoder sizes, number of hidden layers in the encoder, size of the intermediate layers in the encoder etc.\n",
    "\n",
    "* `bert_example/schedulers.py`: contains definitions for learning rate schedulers used in end to end training of the BERT model (`bert_example/sagemaker_smp_pretrain.py`).\n",
    "\n",
    "* `bert_example/utils.py`: This contains different helper utility functions used in end to end training of the BERT model (`bert_example/sagemaker_smp_pretrain.py`).\n",
    "\n",
    "* `bert_example/file_utils.py`: Contains different file utility functions used in model definition (`bert_example/modeling.py`).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Additional Resources\n",
    "If you are a new user of Amazon SageMaker, you may find the following helpful to learn more about SMP and using SageMaker with Pytorch. \n",
    "\n",
    "* To learn more about the SageMaker model parallelism library, see [Model Parallel Distributed Training with SageMaker Distributed](http://docs.aws.amazon.com/sagemaker/latest/dg/model-parallel.html).\n",
    "\n",
    "* To learn more about using the SageMaker Python SDK with Pytorch, see [Using PyTorch with the SageMaker Python SDK](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html).\n",
    "\n",
    "* To learn more about launching a training job in Amazon SageMaker with your own training image, see [Use Your Own Training Algorithms](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html).\n",
    "\n",
    "\n",
    "### Prerequisites \n",
    "\n",
    "1. You must create an S3 bucket to store the input data to be used for training. This bucket must must be located in the same AWS Region you use to launch your training job. This is the AWS Region you use to run this notebook. To learn how, see [Creating a bucket](https://docs.aws.amazon.com/AmazonS3/latest/gsg/CreatingABucket.html) in the Amazon S3 documentation.\n",
    "\n",
    "2. You must download the dataset that you use for training from [Nvidia Deep Learning Examples](https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/LanguageModeling/BERT) and upload it to the S3 bucket you created. To learn more about the datasets and scripts provided to preprocess and download it, see [Getting the data](https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/BERT/README.md#getting-the-data) in the Nvidia Deep Learning Examples repo README. You can also use the [Quick Start Guide](https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/BERT/README.md#quick-start-guide) to learn how to download the dataset. The repository consists of three datasets. Optionally, you can to use the `wiki_only` parameter to only download the Wikipedia dataset. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Amazon SageMaker Initialization\n",
    "\n",
    "Initialize the notebook instance. Get the AWS Region, SageMaker execution role Amazon Resource Name (ARN). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.estimator import Estimator\n",
    "from sagemaker.pytorch import PyTorch\n",
    "import boto3\n",
    "\n",
    "role = get_execution_role() # provide a pre-existing role ARN as an alternative to creating a new role\n",
    "print(f'SageMaker Execution Role:{role}')\n",
    "\n",
    "client = boto3.client('sts')\n",
    "account = client.get_caller_identity()['Account']\n",
    "print(f'AWS account:{account}')\n",
    "\n",
    "session = boto3.session.Session()\n",
    "region = session.region_name\n",
    "print(f'AWS region:{region}')\n",
    "sagemaker_session = sagemaker.session.Session(boto_session=session)\n",
    "import sys\n",
    "print(sys.path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare/Identify your Training Data in Amazon S3\n",
    "\n",
    "If you don't already have the BERT dataset in an S3 bucket, please see the instructions in [Nvidia BERT Example](https://github.com/NVIDIA/DeepLearningExamples/blob/master/PyTorch/LanguageModeling/BERT/README.md) to download the dataset and upload it to a s3 bucket. See the prerequisites at the beginning of this notebook for more information.\n",
    "\n",
    "Uncomment and use the following cell to specify the Amazon S3 bucket and prefix that contains your training data. For example, if your training data is in s3://your-bucket/training, enter `'your-bucket'` for s3_bucket and `'training'` for prefix. Note that your output data will be stored in the same bucket, under the `output/` prefix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_bucket = '<ADD BUCKET>'\n",
    "#prefix = '<ADD PREFIX>'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define SageMaker Data Channels\n",
    "\n",
    "In this step, you define Amazon SageMaker training data channel and output data path. The training data channel identifies where your training data is located in S3. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3train = f's3://{s3_bucket}/{prefix}'\n",
    "train = sagemaker.session.TrainingInput(s3train, distribution='FullyReplicated', \n",
    "                                        s3_data_type='S3Prefix')\n",
    "\n",
    "data_channels = {'train': train}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Set your output data path. This is where model artifacts are stored. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s3_output_location = f's3://{s3_bucket}/output'\n",
    "print(f'your output data will be stored in: s3://{s3_bucket}/output')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define SageMaker Training Job\n",
    "\n",
    "Next, you will use SageMaker Estimator API to define a SageMaker Training Job. You will use a [`PyTorchEstimator`](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/sagemaker.pytorch.html) to define the number and type of EC2 instances Amazon SageMaker uses for training, as well as the size of the volume attached to those instances. \n",
    "\n",
    "You must update the following:\n",
    "* `instance_count`\n",
    "* `instance_type`\n",
    "* `volume_size`\n",
    "\n",
    "See the following sub-sections for more details. \n",
    "\n",
    "### Update the Type and Number of EC2 Instances Used\n",
    "\n",
    "The instance type and number of instances you specify in `instance_type` and `instance_count` respectively will determine the number of GPUs Amazon SageMaker uses during training. Explicitly, `instance_type` will determine the number of GPUs on a single instance and that number will be multiplied by `instance_count`. \n",
    "\n",
    "You must specify values for `instance_type` and `instance_count` so that the total number of GPUs available for training is equal to `partitions` in `config` of `smp.init` in your training script. \n",
    "\n",
    "If you set ddp to `True`, you must ensure that the total number of GPUs available is divisible by `partitions`. The result of the division is inferred to be the number of model replicas to be used for Horovod (data parallelism degree). \n",
    "\n",
    "See [Amazon SageMaker Pricing](https://aws.amazon.com/sagemaker/pricing/) for SageMaker supported instances and cost information. To look up GPUs for each instance types, see [Amazon EC2 Instance Types](https://aws.amazon.com/ec2/instance-types/). Use the section **Accelerated Computing** to see general purpose GPU instances. Note that an ml.p3.2xlarge has the same number of GPUs as an p3.2xlarge.\n",
    "\n",
    "### Update your Volume Size\n",
    "\n",
    "The volume size you specify in `volume_size` must be larger than your input data size.\n",
    "\n",
    "### Set your parameters dictionary for SMP and set custom mpioptions\n",
    "\n",
    "With the parameters dictionary you can configure: the number of microbatches, number of partitions, whether to use data parallelism with ddp, the pipelining strategy, the placement strategy and other BERT specific hyperparameters. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mpi_options = \"-verbose --mca orte_base_help_aggregate 0 \"\n",
    "mpi_options += \"--mca btl_vader_single_copy_mechanism none\"\n",
    "smp_parameters = {\"optimize\": \"speed\", \"microbatches\": 12, \"partitions\": 2, \"ddp\": True, \"pipeline\": \"interleaved\", \"overlapping_allreduce\": True, \"placement_strategy\": \"cluster\", \"memory_weight\": 0.3}\n",
    "timeout = 60 * 60\n",
    "metric_definitions = [{\"Name\": \"base_metric\", \"Regex\": \"<><><><><><>\"}]\n",
    "\n",
    "hyperparameters = {\"input_dir\": \"/opt/ml/input/data/train\",\n",
    "                   \"output_dir\": \"./checkpoints\", \n",
    "                   \"config_file\": \"bert_config.json\", \n",
    "                   \"bert_model\": \"bert-large-uncased\", \n",
    "                   \"train_batch_size\": 48, \n",
    "                   \"max_seq_length\": 128,\n",
    "                   \"max_predictions_per_seq\": 20,\n",
    "                   \"max_steps\": 7038,\n",
    "                   \"warmup_proportion\": 0.2843,\n",
    "                   \"num_steps_per_checkpoint\": 200,\n",
    "                   \"learning_rate\": 6e-3,\n",
    "                   \"seed\": 12439,\n",
    "                   \"steps_this_run\": 500,\n",
    "                   \"allreduce_post_accumulation\": 1,\n",
    "                   \"allreduce_post_accumulation_fp16\": 1,\n",
    "                   \"do_train\": 1,\n",
    "                   \"use_sequential\": 1,\n",
    "                   \"skip_checkpoint\": 1,\n",
    "                   \"smp\": 1,\n",
    "                   \"apply_optimizer\": 1,\n",
    "                   \"json-summary\": \"./dllogger.json\"}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Instantiate Pytorch Estimator with SMP enabled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pytorch_estimator = PyTorch(\"sagemaker_smp_pretrain.py\",\n",
    "                            role=role,\n",
    "                            instance_type=\"ml.p3.16xlarge\",\n",
    "                            volume_size=200,\n",
    "                            instance_count=1,\n",
    "                            sagemaker_session=sagemaker_session,\n",
    "                            py_version=\"py3\",\n",
    "                            framework_version='1.6.0',\n",
    "                            distribution={\n",
    "                                \"smdistributed\": {\n",
    "                                    \"modelparallel\": {\n",
    "                                        \"enabled\": True,\n",
    "                                        \"parameters\": smp_parameters\n",
    "                                    }\n",
    "                                },\n",
    "                                \"mpi\": {\n",
    "                                    \"enabled\": True,\n",
    "                                    \"process_per_host\": 8,\n",
    "                                    \"custom_mpi_options\": mpi_options,\n",
    "                                }\n",
    "                            },\n",
    "                            source_dir='bert_example',\n",
    "                            output_path=s3_output_location,\n",
    "                            max_run=timeout,\n",
    "                            hyperparameters=hyperparameters,\n",
    "                            metric_definitions=metric_definitions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, you will use the estimator to launch the SageMaker training job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pytorch_estimator.fit(data_channels, logs=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_python3",
   "language": "python",
   "name": "conda_python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}